/*
* Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
* Description: Hccl One-Sided Conn
* Author: yanzhenghang
* Create: 2024-09-06
*/

#ifndef HCCL_ONE_SIDED_CONN_H
#define HCCL_ONE_SIDED_CONN_H

#include <hccl/hccl_types.h>
#include <hccl/base.h>
#include "hccl_socket_manager.h"
#include "hccl_network_pub.h"
#include "local_rdma_rma_buffer.h"
#include "local_ipc_rma_buffer.h"
#include "remote_rdma_rma_buffer.h"
#include "remote_ipc_rma_buffer.h"
#include "hccl_one_sided_services.h"
#include "rma_buffer_mgr.h"
#include "notify_pool.h"
#include "transport_roce_mem.h"
#include "transport_ipc_mem.h"
#include "exception_handler.h"
namespace hccl {

class HcclOneSidedConn {
public:
    using LocalRdmaRmaBufferMgr = RmaBufferMgr<BufferKey<uintptr_t, u64>, std::shared_ptr<LocalRdmaRmaBuffer>>;
    using RemoteRdmaRmaBufferMgr = RmaBufferMgr<BufferKey<uintptr_t, u64>, std::shared_ptr<RemoteRdmaRmaBuffer>>;
    using LocalIpcRmaBufferMgr = RmaBufferMgr<BufferKey<uintptr_t, u64>, std::shared_ptr<LocalIpcRmaBuffer>>;
    using RemoteIpcRmaBufferMgr = RmaBufferMgr<BufferKey<uintptr_t, u64>, std::shared_ptr<RemoteIpcRmaBuffer>>;

    struct ProcessInfo {
        u32 pid;
        u32 sdid;
        u32 serverId;
    };

    struct MemDescData {
        u32 localRankId;
        u32 remoteRankId;
        HcclMemDesc memDesc;

        // 从 string 拷贝到 memDesc
        void CopyFromString(const std::string& localMemDescStr) {
            EXCEPTION_THROW_IF_COND_ERR( memcpy_s(&memDesc, sizeof(HcclMemDesc),
                localMemDescStr.c_str(), localMemDescStr.size() + 1) != EOK, "memcpy_s memDesc failed." );
        }

        // 从 memDesc 转换为 string
        std::string CopyToString() const {
            return std::string(reinterpret_cast<const char*>(&memDesc), sizeof(HcclMemDesc) - MEMDESC_OFFSET);
        }
    };

    // 参数超过5个，最终交付前完成优化
    HcclOneSidedConn(const HcclNetDevCtx &netDevCtx, const HcclRankLinkInfo &localRankInfo,
        const HcclRankLinkInfo &remoteRankInfo, std::unique_ptr<HcclSocketManager> &socketManager,
        LocalRdmaRmaBufferMgr &localRdmaRmaBufferMgr, std::unique_ptr<NotifyPool> &notifyPool,
        LocalIpcRmaBufferMgr &localIpcRmaBufferMgr, const HcclDispatcher &dispatcher,
        const std::string &commIdentifier, const bool &useRdma);

    ~HcclOneSidedConn();

    HcclResult Connect(s32 timeoutSec);

    HcclResult ExchangeIpcProcessInfo(const ProcessInfo &localProcess, ProcessInfo &remoteProcess);
    HcclResult ExchangeMemDesc(const HcclMemDescs &localMemDescs, HcclMemDescs &remoteMemDescs, u32 &actualNumOfRemote);

    void EnableMemAccess(const HcclMemDesc &remoteMemDesc, HcclMem &remoteMem);
    void DisableMemAccess(const HcclMemDesc &remoteMemDesc);

    void BatchWrite(const HcclOneSideOpDesc* oneSideDescs, u32 descNum, const rtStream_t& stream);
    void BatchRead(const HcclOneSideOpDesc* oneSideDescs, u32 descNum, const rtStream_t& stream);

private:
    void EnableMemAccessRdma(const HcclMemDesc &remoteMemDesc, HcclMem &remoteMem);
    void EnableMemAccessIpc(const HcclMemDesc &remoteMemDesc, HcclMem &remoteMem);
    void DisableMemAccessRdma(const HcclMemDesc &remoteMemDesc);
    void DisableMemAccessIpc(const HcclMemDesc &remoteMemDesc);

    HcclResult FillRmaBufferSlice(const HcclOneSideOpDesc* desc, RmaBufferSlice& localRmaBufferSlice,
        RmaBufferSlice& remoteRmaBufferSlice);
    HcclResult FillRdmaRmaBufferSlice(const HcclOneSideOpDesc* desc, RmaBufferSlice& localRmaBufferSlice,
        RmaBufferSlice& remoteRmaBufferSlice);
    HcclResult FillIpcRmaBufferSlice(const HcclOneSideOpDesc* desc, RmaBufferSlice& localRmaBufferSlice,
        RmaBufferSlice& remoteRmaBufferSlice);
    HcclResult SendLocalMemDesc(const HcclMemDescs &localMemDescs);
    HcclResult ReceiveRemoteMemDesc(HcclMemDescs &remoteMemDescs, u32 &actualNumOfRemote);

    static constexpr size_t MEMDESC_OFFSET = 2 * sizeof(u32);
    HcclNetDevCtx netDevCtx_{};

    const HcclRankLinkInfo &localRankInfo_;
    HcclRankLinkInfo remoteRankInfo_{};
    std::unique_ptr<HcclSocketManager> &socketManager_;

    LocalRdmaRmaBufferMgr &localRdmaRmaBufferMgr_;
    RemoteRdmaRmaBufferMgr remoteRdmaRmaBufferMgr_{};
    LocalIpcRmaBufferMgr &localIpcRmaBufferMgr_;
    RemoteIpcRmaBufferMgr remoteIpcRmaBufferMgr_{};

    std::shared_ptr<HcclSocket> socket_{};

    std::unique_ptr<NotifyPool> &notifyPool_;

    std::unique_ptr<TransportMem> transportMemPtr_{};
    HcclDispatcher dispatcher_{};

    std::string commIdentifier_;
    bool useRdma_{true};
};
}
#endif